/**
* Copyright 2023 PYSTOR DBS, Nanjing Pengyun Network Technology Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/ 

package py.coordinator.main;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import io.netty.buffer.ByteBuf;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.After;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;
import py.PbRequestResponseHelper;
import py.archive.segment.SegId;
import py.common.DelayManager;
import py.common.RequestIdBuilder;
import py.common.struct.EndPoint;
import py.common.struct.Pair;
import py.coordinator.base.IoContextGenerator;
import py.coordinator.base.NbdRequestResponseGenerator;
import py.coordinator.iorequest.iorequest.IoRequestType;
import py.coordinator.iorequest.iounit.IoUnit;
import py.coordinator.iorequest.iounitcontext.IoUnitContext;
import py.coordinator.iorequest.iounitcontextpacket.IoUnitContextPacket;
import py.coordinator.iorequest.iounitcontextpacket.IoUnitContextPacketImpl;
import py.coordinator.lib.Coordinator;
import py.coordinator.logmanager.ReadIoContextManager;
import py.coordinator.logmanager.ReadMethodCallback;
import py.coordinator.nbd.NbdResponseSender;
import py.coordinator.nbd.ProtocoalConstants;
import py.coordinator.nbd.request.Reply;
import py.coordinator.task.ResendRequest;
import py.coordinator.utils.PbRequestResponsePbHelper;
import py.instance.InstanceId;
import py.membership.IoActionContext;
import py.membership.IoMember;
import py.membership.MemberIoStatus;
import py.membership.SegmentForm;
import py.netty.datanode.PyReadResponse;
import py.netty.exception.ServerProcessException;
import py.netty.exception.TimeoutException;
import py.netty.memory.SimplePooledByteBufAllocator;
import py.proto.Broadcastlog;
import py.proto.Broadcastlog.PbIoUnitResult;
import py.proto.Broadcastlog.PbReadRequestUnit;
import py.proto.Broadcastlog.PbReadResponseUnit;
import py.test.TestBase;
import py.volume.VolumeType;

/**
 * xx.
 */
public class ReadIoContextManagerTest extends TestBase {

  private final Random random = new Random();
  private ReadIoContextManager ioContextManager;
  private SimplePooledByteBufAllocator allocator = new SimplePooledByteBufAllocator(
      16 * 1024 * 1024, 4 * 1024);
  private DelayManager delayResendManager = mock(DelayManager.class);
  private Coordinator coordinator = mock(Coordinator.class);
  private EndPoint endPoint = mock(EndPoint.class);
  private SegId segId = mock(SegId.class);
  private IoContextGenerator ioContextGenerator;
  private long segmentSize = 1024 * 1024;
  private int pageSize = 128;
  private int pageCount;
  private int littlePageCount;
  private List<Reply> replies = new ArrayList<>();
  private IoActionContext ioActionContext = mock(IoActionContext.class);
  private Long volumeId = RequestIdBuilder.get();


  /**
   * xx.
   */
  public ReadIoContextManagerTest() throws Exception {
    super.init();
    ioContextGenerator = new IoContextGenerator(volumeId, segmentSize, pageSize);
    when(coordinator.getDelayManager()).thenReturn(delayResendManager);
    when(coordinator.getVolumeType(any(Long.class))).thenReturn(VolumeType.REGULAR);
    pageCount = allocator.getAvailableMediumPageCount();
    littlePageCount = allocator.getAvailableLittlePageCount();
  }


  /**
   * xx.
   */
  @Before
  public void beforeMethod() {
    // TODO: SegmentForm.PJ
    when(ioActionContext.getSegmentForm()).thenReturn(SegmentForm.PJ);
    Set<IoMember> ioMembers = new HashSet<>();
    ioMembers
        .add(new IoMember(new InstanceId(RequestIdBuilder.get()), new EndPoint("127.0.0.1", 6789),
            MemberIoStatus.Primary, Broadcastlog.ReadCause.FETCH));
    when(ioActionContext.getIoMembers()).thenReturn(ioMembers);
    when(ioActionContext.getRealReaders()).thenReturn(ioMembers);
    when(ioActionContext.getFetchReader()).thenReturn(ioMembers);
    replies.clear();
  }

  @After
  public void afterMethod() {
    assertEquals(pageCount, allocator.getAvailableMediumPageCount());
    assertEquals(littlePageCount, allocator.getAvailableLittlePageCount());
  }

  /**
   * stimulate one read request with good response.
   */
  @Test
  public void oneGoodReadResponse() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        logger.warn("reply request: {}", reply.getResponse());
        replies.add(reply);
        NbdRequestResponseGenerator.checkBuffer(reply.getResponse().getBody(), 1);
        if (reply.getResponse().getBody() != null) {
          reply.getResponse().getBody().release();
        }
        assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.SUCCEEDED);
      }
    };
    ioContextGenerator.setSender(sender);

    List<IoUnitContext> contexts = createReadIoContextManager(volumeId, 127, 2 * pageSize);
    assertEquals(contexts.size(), 3);

    responseGood();

    assertEquals(replies.size(), 1);
    assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.SUCCEEDED);
    verify(delayResendManager, times(0)).put(any(ResendRequest.class));
  }

  /**
   * stimulate one read request with two good responses, cause first round bytebuf has been
   * released.
   */
  @Ignore
  @Test
  public void twoGoodReadResponse() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        replies.add(reply);
        NbdRequestResponseGenerator.checkBuffer(reply.getResponse().getBody(), 1);
        if (reply.getResponse().getBody() != null) {
          reply.getResponse().getBody().release();
        }
        assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.SUCCEEDED);
      }
    };
    ioContextGenerator.setSender(sender);
    List<IoUnitContext> totalContexts = createReadIoContextManager(volumeId, 125, 4 * pageSize);
    assertEquals(totalContexts.size(), 5);
    List<IoUnitContext> contexts = null;

    // submit the first 3 contexts
    contexts = totalContexts.subList(0, 3);
    ioContextManager = new ReadIoContextManager(volumeId, segId,
        new IoUnitContextPacketImpl(volumeId, contexts, 0, IoRequestType.Read),
        coordinator);
    ioContextManager.setIoActionContext(ioActionContext);
    responseGood();
    assertEquals(replies.size(), 0);
    assertEquals(totalContexts.get(0).getIoRequest().getReferenceCount(), 2);

    // submit the next 2 contexts
    contexts = totalContexts;
    ioContextManager = new ReadIoContextManager(volumeId, segId,
        new IoUnitContextPacketImpl(volumeId, contexts, 0, IoRequestType.Read),
        coordinator);
    ioContextManager.setIoActionContext(ioActionContext);
    responseGood();

    assertEquals(replies.size(), 1);

    verify(delayResendManager, times(0)).put(any(ResendRequest.class));
  }

  /**
   * meet timeout and fail to response the request.
   */
  @Test
  public void timeoutAndFailResponse() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        logger.info("catch reply: {}", reply);
        replies.add(reply);
        if (reply.getResponse().getBody() != null) {
          reply.getResponse().getBody().release();
        }

        assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.EIO);
      }
    };
    ioContextGenerator.setSender(sender);
    List<IoUnitContext> contexts = createReadIoContextManager(volumeId, 127, 2 * pageSize);

    assertTrue(contexts.size() == 3);
    // retry

    responseFail(new TimeoutException());
    assertEquals(replies.size(), 0);
    verify(delayResendManager, times(1)).put(any(ResendRequest.class));

    // timeout and response
    ioContextManager.setExpiredTime(System.currentTimeMillis());
    Thread.sleep(1);
    responseFail(new TimeoutException());

    assertEquals(replies.size(), 1);
    verify(delayResendManager, times(1)).put(any(ResendRequest.class));
  }

  /**
   * a request is splited as many request units, some units will be responsed as ok, others will be
   * responsed as failure. the request will be responsed as failure.
   */
  @Test
  public void someGoodAndSomeBadResponse() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        replies.add(reply);
        if (reply.getResponse().getBody() != null) {
          logger.warn("releaseBody: {}", reply);
          reply.getResponse().getBody().release();
        }

        assertEquals(reply.getResponse().getErrCode(), ProtocoalConstants.EIO);
      }
    };

    ioContextGenerator.setSender(sender);
    List<IoUnitContext> totalContexts = createReadIoContextManager(volumeId, 125, 4 * pageSize);
    assertEquals(totalContexts.size(), 5);
    // response three good, two not good.
    int successCount = 3;
    responseNotAllGood(successCount, PbIoUnitResult.SKIP);
    verify(delayResendManager, times(1)).put(any(ResendRequest.class));
    assertEquals(replies.size(), 0);
    assertEquals(totalContexts.get(3).getIoRequest().getReferenceCount(), 5 - 3);
    for (int i = 0; i < totalContexts.size(); i++) {
      if (i < successCount) {
        assertTrue(totalContexts.get(i).getIoUnit().isSuccess());
      } else {
        assertFalse(totalContexts.get(i).getIoUnit().isSuccess());
      }
    }

    // response not good, and expired.
    ioContextManager.setExpiredTime(System.currentTimeMillis());
    Thread.sleep(1);
    responseFail(new TimeoutException());
    verify(delayResendManager, times(1)).put(any(ResendRequest.class));
    assertEquals(replies.size(), 1);
  }

  /**
   * the result of the request unit will be outofrange or input no data.
   */
  @Test
  public void badResponse() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        replies.add(reply);
        assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.EIO);
      }
    };
    ioContextGenerator.setSender(sender);
    List<IoUnitContext> contexts = createReadIoContextManager(volumeId, 127, 2 * pageSize);
    assertTrue(contexts.size() == 3);

    // first time to response fail.
    responseFail(new ServerProcessException("one exception"));
    verify(delayResendManager, times(1)).put(any(ResendRequest.class));

    // second time to response fail.
    responseFail(new TimeoutException("two exception"));
    verify(delayResendManager, times(2)).put(any(ResendRequest.class));
    assertEquals(replies.size(), 0);

    // third time to response fail.
    responseNotAllGood(1,
        random.nextBoolean() ? PbIoUnitResult.OUT_OF_RANGE : PbIoUnitResult.INPUT_HAS_NO_DATA);

    assertEquals(replies.size(), 1);
    verify(delayResendManager, times(2)).put(any(ResendRequest.class));
  }

  /**
   * resend until the request is success.
   */
  @Test
  public void resendUntilSuccess() throws Exception {
    NbdResponseSender sender = new NbdResponseSender() {
      @Override
      public void send(Reply reply) {
        replies.add(reply);
        if (reply.getResponse().getBody() != null) {
          reply.getResponse().getBody().release();
        }
        assertEquals(replies.get(0).getResponse().getErrCode(), ProtocoalConstants.SUCCEEDED);
      }
    };
    ioContextGenerator.setSender(sender);
    List<IoUnitContext> contexts = createReadIoContextManager(volumeId, 127, 2 * pageSize);
    logger.info("contexts : {}", contexts.size());
    assertTrue(contexts.size() == 3);
    int retryTimes = 3;
    for (int i = 0; i < retryTimes; i++) {
      responseFail(new TimeoutException());
      verify(delayResendManager, times(i + 1)).put(any(ResendRequest.class));
    }

    assertEquals(replies.size(), 0);
    responseGood();

    assertEquals(replies.size(), 1);
    verify(delayResendManager, times(retryTimes)).put(any(ResendRequest.class));

  }

  private List<IoUnitContext> createReadIoContextManager(long volumeId, long offset, int length)
      throws Exception {
    List<IoUnitContext> contexts = ioContextGenerator
        .generateReadIoContexts(volumeId, offset, length);
    IoUnitContextPacket ioUnitContextPacket = new IoUnitContextPacketImpl(volumeId, contexts, 0,
        IoRequestType.Read);
    ioContextManager = new ReadIoContextManager(volumeId, segId, ioUnitContextPacket, coordinator);
    ioContextManager.setExpiredTime(System.currentTimeMillis() + 10000000);
    ioContextManager.setIoActionContext(ioActionContext);
    return contexts;
  }

  private void responseGood() {
    IoMember ioMember = new IoMember(new InstanceId(RequestIdBuilder.get()), endPoint,
        MemberIoStatus.Primary,
        Broadcastlog.ReadCause.FETCH);
    // TODO: send count should change
    ReadMethodCallback callback = new ReadMethodCallback(ioContextManager, new AtomicInteger(1),
        coordinator,
        ioMember);
    ReadMethodCallback[] callbacks = new ReadMethodCallback[1];
    callbacks[0] = callback;
    ioContextManager.initRequestCount(callbacks.length);
    Pair<PbReadResponseUnit, ByteBuf>[] pairs = generateContext(ioContextManager.getRequestUnits(),
        PbIoUnitResult.OK);
    PyReadResponse response = NbdRequestResponseGenerator
        .generateReadResponse(RequestIdBuilder.get(), endPoint, pairs);
    callback.complete(response);
  }

  private void responseFail(Exception e) {
    IoMember ioMember = new IoMember(new InstanceId(RequestIdBuilder.get()), endPoint,
        MemberIoStatus.Primary,
        Broadcastlog.ReadCause.FETCH);
    ReadMethodCallback callback = new ReadMethodCallback(ioContextManager, new AtomicInteger(1),
        coordinator,
        ioMember);
    ReadMethodCallback[] callbacks = new ReadMethodCallback[1];
    callbacks[0] = callback;
    ioContextManager.initRequestCount(callbacks.length);
    callback.fail(e);
  }

  @SuppressWarnings("unchecked")
  private void responseNotAllGood(int successCount, PbIoUnitResult result) {
    List<IoUnitContext> ioContexts = ioContextManager.getRequestUnits();
    IoMember ioMember = new IoMember(new InstanceId(RequestIdBuilder.get()), endPoint,
        MemberIoStatus.Primary,
        Broadcastlog.ReadCause.FETCH);
    ReadMethodCallback callback = new ReadMethodCallback(ioContextManager, new AtomicInteger(1),
        coordinator,
        ioMember);
    ReadMethodCallback[] callbacks = new ReadMethodCallback[1];
    callbacks[0] = callback;
    ioContextManager.initRequestCount(callbacks.length);
    Pair<PbReadResponseUnit, ByteBuf>[] pairs = new Pair[ioContexts.size()];
    Pair<PbReadResponseUnit, ByteBuf>[] tmp = generateContext(ioContexts.subList(0, successCount),
        PbIoUnitResult.OK);
    for (int i = 0; i < successCount; i++) {
      pairs[i] = tmp[i];
    }

    tmp = generateContext(ioContexts.subList(successCount, ioContexts.size()), result);
    for (int i = 0; i < ioContexts.size() - successCount; i++) {
      pairs[successCount + i] = tmp[i];
    }

    PyReadResponse response = NbdRequestResponseGenerator
        .generateReadResponse(RequestIdBuilder.get(), endPoint, pairs);
    callback.complete(response);
  }


  /**
   * xx.
   */
  @SuppressWarnings({"unchecked", "rawtypes"})
  public Pair<PbReadResponseUnit, ByteBuf>[] generateContext(List<IoUnitContext> contexts,
      PbIoUnitResult result) {
    Pair<PbReadResponseUnit, ByteBuf>[] pairs = new Pair[contexts.size()];
    for (int i = 0; i < contexts.size(); i++) {
      IoUnit ioUnit = contexts.get(i).getIoUnit();
      PbReadRequestUnit unit = PbRequestResponsePbHelper.buildPbReadRequestUnitFrom(ioUnit);
      if (result == PbIoUnitResult.OK) {

        byte[] data = NbdRequestResponseGenerator.getBuffer(ioUnit.getLength(), 1);
        ByteBuf buf = allocator.buffer(ioUnit.getLength());
        buf.writeBytes(data);
        PbReadResponseUnit responseUnit = PbRequestResponseHelper
            .buildPbReadResponseUnitFrom(unit, buf);
        pairs[i] = new Pair(responseUnit, buf);
      } else {
        PbReadResponseUnit responseUnit = PbRequestResponseHelper
            .buildPbReadResponseUnitFrom(unit, result);
        pairs[i] = new Pair(responseUnit, null);
      }
    }

    return pairs;
  }
}
